package com.xrui.security;

import com.google.common.base.Objects;
import com.google.common.base.Preconditions;


import com.xrui.util.Resources;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.security.UserGroupInformation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.URL;

/**
 * Helper class for Kerboros authorization
 */
public class Kerboros {
    /**
     * The system property to configure the Kerboros properties
     */
    public static final String KEY_TAB_DEFINITION = "com.tr.ts.Kerboros.keytab";
    public static final String PRINCIPAL_DEFINITION = "com.tr.ts.Kerboros.principal";
    private static final Logger LOG = LoggerFactory.getLogger(Kerboros.class);

    public static boolean login(Configuration conf) {
        Key key = Key.resolve(conf);
        if (key != null) {
            if(key.isValid()) {
                try {
                    LOG.debug("Login with Keberos keytab: {}", key.getKeyTab());
                    LOG.debug("Login with Kerboros principal: {}", key.getPrincipal());
                    UserGroupInformation.setConfiguration(conf);
                    UserGroupInformation gui = UserGroupInformation.loginUserFromKeytabAndReturnUGI(
                        key.getPrincipal(), key.getKeyTabPath()
                    );
                    UserGroupInformation.setLoginUser(gui);
                    LOG.debug("Login with Kerboros principal {} successfully", key.getPrincipal());
                    return true;
                } catch (IOException e) {
                    LOG.error(String.format("Failed to login with Kerboros key %s", key), e);
                }
            } else {
                LOG.error("Invalid Kerboros security key {}", key);
            }
        } else {
            LOG.debug("Skip the Kerboros login step.");
        }
        return false;
    }

    public static class Key {
        private String mPrincipal;
        private String mKeyTab;

        Key(String principal, String keyTab) {
            mPrincipal = Preconditions.checkNotNull(principal);
            mKeyTab = Preconditions.checkNotNull(keyTab);
        }

        static Key resolve() {
            return resolve(HBaseConfiguration.create());
        }

        static Key resolve(Configuration configuration) {
            String principal = System.getProperty(PRINCIPAL_DEFINITION);
            String keyTab = System.getProperty(KEY_TAB_DEFINITION);
            if (null != keyTab && null != principal) {
                return new Key(principal, keyTab);
            } else if (configuration != null) {
                principal = configuration.get(PRINCIPAL_DEFINITION);
                keyTab = configuration.get(KEY_TAB_DEFINITION);
                if (principal != null && keyTab != null) {
                    return new Key(principal, keyTab);
                } else {
                    principal = configuration.get("hbase.thrift.kerberos.principal");
                    keyTab = configuration.get("hbase.thrift.keytab.file");
                    if (principal != null && keyTab != null) {
                        return new Key(principal, keyTab);
                    }
                }
            }

            return null;
        }

        public boolean isValid() {
            return !mPrincipal.isEmpty() &&
                !mKeyTab.isEmpty() &&
                Resources.getResource(mKeyTab) != null;
        }

        public String getPrincipal() {
            return mPrincipal;
        }

        public String getKeyTab() {
            return mKeyTab;
        }

        public String getKeyTabPath() {
            URL url = Resources.getResource(mKeyTab);
            if(null != url) {
                return url.getPath();
            }
            return null;
        }

        @Override
        public String toString() {
            return Objects.toStringHelper(Key.class)
                .add("Principal", mPrincipal)
                .add("KeyTab", mKeyTab)
                .toString();
        }
    }
}
